import pygame
import random

pygame.init()

screen_width = 1024
screen_height = 768
screen = pygame.display.set_mode((screen_width, screen_height))

# 颜色定义
colors = [
    (0, 0, 0),  # 黑色背景
    (255, 0, 0), (0, 255, 0), (0, 0, 255),
    (255, 255, 0), (255, 0, 255), (0, 255, 255), (128, 128, 128)
]

# 定义所有可能的形状
shapes = [
    [[1, 1, 1, 1]],  # I 形
    [[2, 2], [2, 2]],  # O 形
    [[0, 3, 3], [3, 3, 0]],  # S 形
    [[4, 4, 0], [0, 4, 4]],  # Z 形
    [[0, 5, 0], [5, 5, 5]],  # T 形
    [[6, 0, 0], [6, 6, 6]],  # J 形
    [[0, 0, 7], [7, 7, 7]]  # L 形
]

# 定义形状到颜色索引的映射
shape_to_color_index = {
    1: 1, 2: 2, 3: 3, 4: 4, 5: 5, 6: 6, 7: 7
}

def transpose(matrix):
    return [[matrix[j][i] for j in range(len(matrix))] for i in range(len(matrix[0]))]

def mirror_horizontal(matrix):
    return [row[::-1] for row in matrix]

def mirror_vertical(matrix):
    return matrix[::-1]

def find_all_shapes(shape):
    if shape==mirror_horizontal(transpose(shape)):
        return [shape]
    elif shape == mirror_vertical(mirror_horizontal(shape)):
        return [shape,mirror_horizontal(transpose(shape))]
    else:
        return [shape,mirror_horizontal(transpose(shape)),mirror_vertical(mirror_horizontal(shape)),transpose(shape)]

# 方块类
class Piece(object):
    def __init__(self, x, y, startshape):
        self.x = x
        self.y = y
        self.shape = find_all_shapes(startshape)
        # 找到形状中的第一个非零值来确定颜色
        print(startshape)
        for row in startshape:
            for cell in row:
                if cell != 0:
                    self.color = colors[shape_to_color_index[cell]]
                    break
            if self.color != (0, 0, 0):
                break
        self.rotation = 0

# 创建游戏网格
def create_grid(locked_positions={}):
    grid = [[(0, 0, 0) for _ in range(grid_width)] for _ in range(grid_height)]
    for i in range(grid_height):
        for j in range(grid_width):
            if (j, i) in locked_positions:
                c = locked_positions[(j, i)]
                grid[i][j] = c
    return grid

# 转换形状位置
def convert_shape_format(piece):
    positions = []
    format = piece.shape[piece.rotation % len(piece.shape)]

    for i, line in enumerate(format):
        row = list(line)
        for j, column in enumerate(row):
            if column != 0:  # 过滤掉值为0的位置
                positions.append((piece.x + j, piece.y + i))

    for i, pos in enumerate(positions):
        positions[i] = (pos[0] - 2, pos[1] - 4)

    return positions

# 检查是否超出边界或碰撞
def valid_space(piece, grid):
    accepted_pos = [[(j, i) for j in range(grid_width) if grid[i][j] == (0, 0, 0)] for i in range(grid_height)]
    accepted_pos = [j for sub in accepted_pos for j in sub]
    
    formatted = convert_shape_format(piece)
    
    for pos in formatted:
        if pos not in accepted_pos:
            if pos[1] > -1:
                return False
    return True

# 检查是否输掉游戏
def check_lost(positions):
    for pos in positions:
        x, y = pos
        if y < 1:
            return True
    return False

# 获取新形状
def get_shape():
    return Piece(5, 0, random.choice(shapes))

# 绘制文本中间
def draw_text_middle(text, size, color, surface):
    font = pygame.font.SysFont('comicsans', size, bold=True)
    label = font.render(text, 1, color)
    
    surface.blit(label, (top_left_x + play_width/2 - (label.get_width() / 2), top_left_y + play_height/2 - label.get_height()/2))

# 绘制网格
def draw_grid(surface, grid):
    sx = top_left_x
    sy = top_left_y
    
    for i in range(grid_height):
        pygame.draw.line(surface, (128, 128, 128), (sx, sy + i * grid_size), (sx + play_width, sy + i * grid_size))
        for j in range(grid_width):
            pygame.draw.line(surface, (128, 128, 128), (sx + j * grid_size, sy), (sx + j * grid_size, sy + play_height))

# 清除行
def clear_rows(grid, locked):
    inc = 0
    for i in range(grid_height - 1, -1, -1):
        row = grid[i]
        if (0, 0, 0) not in row:
            inc += 1
            ind = i
            for j in range(grid_width):
                try:
                    del locked[(j, i)]
                except:
                    continue
    
    if inc > 0:
        for key in sorted(list(locked), key=lambda x: x[1])[::-1]:
            x, y = key
            if y < ind:
                newKey = (x, y + inc)
                locked[newKey] = locked.pop(key)
    
    return inc

# 绘制下一个形状
def draw_next_shape(piece, surface):
    font = pygame.font.SysFont('comicsans', 30)
    label = font.render('Next Shape', 1, (255, 255, 255))
    
    sx = top_left_x + play_width + 50
    sy = top_left_y + play_height/2 - 100
    format = piece.shape[piece.rotation % len(piece.shape)]
    
    for i, line in enumerate(format):
        row = list(line)
        for j, column in enumerate(row):
            if column != 0:
                pygame.draw.rect(surface, piece.color, (sx + j*grid_size, sy + i*grid_size, grid_size, grid_size), 0)
    
    surface.blit(label, (sx + 10, sy - 30))

# 更新分数
def update_score(nscore):
    score = max_score()
    
    with open('scores.txt', 'w') as f:
        if int(score) > nscore:
            f.write(str(score))
        else:
            f.write(str(nscore))

def max_score():
    try:
        with open('scores.txt', 'r') as f:
            lines = f.readlines()
            score = lines[0].strip()
    except FileNotFoundError:
        score = '0'
    
    return score

# 绘制窗口
def draw_window(surface, grid):
    surface.fill((0, 0, 0))
    
    pygame.font.init()
    font = pygame.font.SysFont('comicsans', 60)
    label = font.render('Tetris', 1, (255, 255, 255))
    
    surface.blit(label, (top_left_x + play_width / 2 - (label.get_width() / 2), 30))
    
    for i in range(grid_height):
        for j in range(grid_width):
            pygame.draw.rect(surface, grid[i][j], (top_left_x + j*grid_size, top_left_y + i*grid_size, grid_size, grid_size), 0)
    
    draw_grid(surface, grid)
    pygame.draw.rect(surface, (255, 0, 0), (top_left_x, top_left_y, play_width, play_height), 5)

# 主游戏循环
def main():
    global grid

    locked_positions = {}
    grid = create_grid(locked_positions)

    change_piece = False
    run = True
    current_piece = get_shape()
    next_piece = get_shape()
    clock = pygame.time.Clock()
    fall_time = 0
    score = 0

    while run:
        grid = create_grid(locked_positions)
        fall_speed = 0.27
        
        fall_time += clock.get_rawtime()
        clock.tick()
        
        if fall_time/1000 >= fall_speed:
            fall_time = 0
            current_piece.y += 1
            if not valid_space(current_piece, grid) and current_piece.y > 0:
                current_piece.y -= 1
                change_piece = True

        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                run = False
                pygame.display.quit()
                quit()
                
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_LEFT:
                    current_piece.x -= 1
                    if not valid_space(current_piece, grid):
                        current_piece.x += 1
                elif event.key == pygame.K_RIGHT:
                    current_piece.x += 1
                    if not valid_space(current_piece, grid):
                        current_piece.x -= 1
                elif event.key == pygame.K_DOWN:
                    current_piece.y += 1
                    if not valid_space(current_piece, grid):
                        current_piece.y -= 1
                elif event.key == pygame.K_UP:
                    current_piece.rotation = (current_piece.rotation + 1) % len(current_piece.shape)
                    if not valid_space(current_piece, grid):
                        current_piece.rotation = (current_piece.rotation - 1) % len(current_piece.shape)
        
        shape_pos = convert_shape_format(current_piece)
        
        for i in range(len(shape_pos)):
            x, y = shape_pos[i]
            if y > -1:
                grid[y][x] = current_piece.color
        
        if change_piece:
            for pos in shape_pos:
                p = (pos[0], pos[1])
                locked_positions[p] = current_piece.color
            current_piece = next_piece
            next_piece = get_shape()
            change_piece = False
            
            if check_lost(locked_positions):
                run = False
                
        draw_window(screen, grid)
        draw_next_shape(next_piece, screen)
        pygame.display.update()
        
        if check_lost(locked_positions):
            draw_text_middle("YOU LOST!", 80, (255, 255, 255), screen)
            pygame.display.update()
            pygame.time.delay(1500)
            run = False
            update_score(score)

# 主菜单
def main_menu():
    run = True
    while run:
        screen.fill((0, 0, 0))
        draw_text_middle('Press Enter to Begin', 60, (255, 255, 255), screen)
        pygame.display.update()
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                run = False
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_RETURN:
                    main()
    pygame.quit()

# 主菜单
top_left_x = 300
top_left_y = 100
play_width = 300
play_height = 600
grid_size = 30
grid_width = play_width // grid_size
grid_height = play_height // grid_size

main_menu()